
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include "../iface/debug.h"
#include "../utils/BaseConvert.h"
#include "../utils/StringUtils.h"
#include "ECDSAPublicKey.h"

/** Constructor for a blank ECDSA public key. */
ECDSAPublicKey::ECDSAPublicKey()
{
   key = NULL;
   keyname = NULL;
}

/** Constructor for a blank ECDSA public key with a name.
 * 
 * @param name
 * The name to attach to the blank key. */
ECDSAPublicKey::ECDSAPublicKey(const char *name)
{
   key = NULL;
   keyname = NULL;
   setKeyName(name);
}

/** Copy constructor for an ECDSA public key.
 * 
 * @param other
 * The ECDSA public key to copy from. */
ECDSAPublicKey::ECDSAPublicKey(const ECDSAPublicKey &other)
{
   key = EC_KEY_new();
   if (key == NULL)
   {
      errprint("EC_KEY_new failed");
      return;
   }
   if (copyECDSASSLPublicKey(key, other.key))
   {
      errprint("EC_KEY_new failed");
      return;
   }
   keyname = NULL;
   setKeyName(other.keyname);
}

/** Construct an ECDSA public key from an ECDSA public key in PEM encoded 
 * PKCS #8 text format.
 * 
 * @param pemtext
 * The PEM encoded PKCS #8 text that represents an ECDSA public key, from
 * which this key is to be constructed.
 * 
 * @param name
 * The name to be attached to the key, once constructed. */
ECDSAPublicKey::ECDSAPublicKey(const char *pemtext, const char *name)
{
   importKeyFromASCII(pemtext, name);
}

/** Construct an ECDSA public key from an ECDSA public key in SSL format.
 * 
 * @param pubkey
 * The SSL format ECDSA public key to construct this key from.
 * 
 * @param name
 * The name to be attached to the key, once constructed. */
ECDSAPublicKey::ECDSAPublicKey(EC_KEY *pubkey, const char *name)
{
   key = EC_KEY_new();
   if (key == NULL)
   {
      errprint("EC_KEY_new failed");
      return;
   }
   int errval = copyECDSASSLPublicKey(key, pubkey);
   keyname = NULL;
   setKeyName(name);
}

/** Copy constructor for an ECDSA public key.
 * 
 * @param other
 * The ECDSA public key to copy from.
 * 
 * @return
 * Itself. */
ECDSAPublicKey &ECDSAPublicKey::operator=(const ECDSAPublicKey &other)
{
   key = EC_KEY_new();
   if (key == NULL)
   {
      errprint("EC_KEY_new failed");
      return *this;
   }
   if (copyECDSASSLPublicKey(key, other.key))
   {
      errprint("EC_KEY_new failed");
      return *this;
   }
   keyname = NULL;
   setKeyName(other.keyname);
   return *this;
}

/** Destructor for the ECDSA public key. */
ECDSAPublicKey::~ECDSAPublicKey()
{
   if (key) EC_KEY_free(key);
   key = NULL;
   if (keyname) free((void *)keyname);
   keyname = NULL;
}

/** Obtains a cryptlib context which has had the key loaded into it.
 * 
 * @param[out] keycontext
 * A pointer to a cryptlib CRYPT_CONTEXT which is to be initialised with the
 * key.
 * 
 * @return
 * Zero on success, non-zero on failure. */
int ECDSAPublicKey::getCryptlibContext(CRYPT_CONTEXT *keycontext)
{
   int status;
   CRYPT_PKCINFO_ECC *ecdsakey = (CRYPT_PKCINFO_ECC *)malloc(sizeof(CRYPT_PKCINFO_ECC));
   if (getCryptlibKeyComponents(ecdsakey) != 0)
   {
      free(ecdsakey);
      return -1;
   }
   int magictriescntxt = 0;
   do
   {
      status = cryptCreateContext(keycontext, CRYPT_UNUSED, CRYPT_ALGO_ECDSA);
      if (status == CRYPT_ERROR_TIMEOUT) apr_sleep(rand() % CRYPT_MAX_TIMESLEEP);
      magictriescntxt++;
   } while ((status == CRYPT_ERROR_TIMEOUT) && (magictriescntxt < CRYPT_MAX_RETRIES));
   fCryptCheckError(status);
   if (status != CRYPT_OK) goto failover;
   status = cryptSetAttributeString(*keycontext, CRYPT_CTXINFO_LABEL, keyname, strlen(keyname));
   fCryptCheckError(status);
   if (status != CRYPT_OK) goto failover;
   status = cryptSetAttributeString(*keycontext, CRYPT_CTXINFO_KEY_COMPONENTS, ecdsakey, sizeof(CRYPT_PKCINFO_ECC));
   fCryptCheckError(status);
   if (status != CRYPT_OK) goto failover;
   cryptDestroyComponents(ecdsakey);
   free(ecdsakey);
   return 0;
failover:
   cryptDestroyComponents(ecdsakey);
   free(ecdsakey);
   return -1;
}

/** Obtains a cryptlib context with key meant particularly for signature
 * verification.
 * 
 * @param[out] keycontext
 * A pointer to a cryptlib CRYPT_CONTEXT which is to be initialised with the
 * key.
 * 
 * @return
 * Zero on success, non-zero on failure. */
int ECDSAPublicKey::getCryptlibCheckingContext(CRYPT_CONTEXT *keycontext)
{
   return getCryptlibContext(keycontext);
}

/** Obtains a set of key components that can be loaded into cryptlib as a
 * public key.
 * 
 * @param[out] ecdsakey
 * A pointer to a cryptlib ECDSA key components structure which is to be
 * initialised with the key.
 * 
 * @return
 * Zero on success, non-zero on failure. */
int ECDSAPublicKey::getCryptlibKeyComponents(CRYPT_PKCINFO_ECC *ecdsakey)
{
   BN_CTX *ctx = BN_CTX_new();
   BIGNUM *p = BN_new();
   BIGNUM *a = BN_new();
   BIGNUM *b = BN_new();
   BIGNUM *gx = BN_new();
   BIGNUM *gy = BN_new();
   BIGNUM *n = BN_new();
   BIGNUM *h = BN_new();
   BIGNUM *qx = BN_new();
   BIGNUM *qy = BN_new();
   unsigned char *buffera[9] = { NULL, NULL, NULL, NULL, NULL,
                                 NULL, NULL, NULL, NULL }; // p, a, b, gx, gy, n, h, qx, qy
   int failflag = 0;
   int curveret = 0, genpointret = 0, getorder = 0, getcofactor = 0, pubkeyret = 0;
   const EC_POINT *ecpointret, *ecpublickey;
   if ((!ctx) || (!p) || (!a) || (!b) || (!gx) || (!gy) ||
       (!n) || (!h) || (!qx) || (!qy))
      { failflag = 1; errprint("%s creation failed", (ctx) ? "bignum context" : "bignum"); goto cleanup; }
   curveret = EC_GROUP_get_curve_GFp(EC_KEY_get0_group(key),
                                         p, a, b, ctx);
   if (!curveret) { failflag = 1; errprint("could not get curve parameters"); goto cleanup; }
   ecpointret = EC_GROUP_get0_generator(EC_KEY_get0_group(key));
   if (!ecpointret) { failflag = 1; errprint("could not get generator point"); goto cleanup; }
   genpointret = EC_POINT_get_affine_coordinates_GFp(EC_KEY_get0_group(key),
                                                         ecpointret, gx, gy, ctx);
   if (!genpointret) { failflag = 1; errprint("could not get generator point coordinates"); goto cleanup; }
   getorder = EC_GROUP_get_order(EC_KEY_get0_group(key),
                                     n, ctx);
   if (!getorder) { failflag = 1; errprint("could not get order parameter"); goto cleanup; }
   getcofactor = EC_GROUP_get_cofactor(EC_KEY_get0_group(key),
                                           h, ctx);
   if (!getcofactor) { failflag = 1; errprint("could not get cofactor parameter"); goto cleanup; }
   ecpublickey = EC_KEY_get0_public_key(key);
   if (!ecpublickey) { failflag = 1; errprint("could not get public key"); goto cleanup; }
   pubkeyret = EC_POINT_get_affine_coordinates_GFp(EC_KEY_get0_group(key),
                                                       ecpublickey, qx, qy, ctx);
   if (!ecpublickey) { failflag = 1; errprint("could not get public point coordinates"); goto cleanup; }

   cryptInitComponents(ecdsakey, CRYPT_KEYTYPE_PUBLIC);   
   buffera[0] = (unsigned char *)malloc(BN_num_bytes(p));
   if (!(buffera[0])) { errprint("Could not allocate buffer for key->p."); return -1; }
   buffera[1] = (unsigned char *)malloc(BN_num_bytes(a));
   if (!(buffera[1])) { errprint("Could not allocate buffer for key->a."); return -1; }
   buffera[2] = (unsigned char *)malloc(BN_num_bytes(b));
   if (!(buffera[2])) { errprint("Could not allocate buffer for key->b."); return -1; }
   buffera[3] = (unsigned char *)malloc(BN_num_bytes(gx));
   if (!(buffera[3])) { errprint("Could not allocate buffer for key->gx."); return -1; }
   buffera[4] = (unsigned char *)malloc(BN_num_bytes(gy));
   if (!(buffera[4])) { errprint("Could not allocate buffer for key->gy."); return -1; }
   buffera[5] = (unsigned char *)malloc(BN_num_bytes(n));
   if (!(buffera[5])) { errprint("Could not allocate buffer for key->n."); return -1; }
   buffera[6] = (unsigned char *)malloc(BN_num_bytes(h));
   if (!(buffera[6])) { errprint("Could not allocate buffer for key->h."); return -1; }
   buffera[7] = (unsigned char *)malloc(BN_num_bytes(qx));
   if (!(buffera[7])) { errprint("Could not allocate buffer for key->qx."); return -1; }
   buffera[8] = (unsigned char *)malloc(BN_num_bytes(qy));
   if (!(buffera[8])) { errprint("Could not allocate buffer for key->qy."); return -1; }

   BN_bn2bin(p, buffera[0]);
   BN_bn2bin(a, buffera[1]);
   BN_bn2bin(b, buffera[2]);
   BN_bn2bin(gx, buffera[3]);
   BN_bn2bin(gy, buffera[4]);
   BN_bn2bin(n, buffera[5]);
   BN_bn2bin(h, buffera[6]);
   BN_bn2bin(qx, buffera[7]);
   BN_bn2bin(qy, buffera[8]);
   cryptSetComponent(ecdsakey->p, buffera[0], BN_num_bits(p));
   cryptSetComponent(ecdsakey->a, buffera[1], BN_num_bits(a));
   cryptSetComponent(ecdsakey->b, buffera[2], BN_num_bits(b));
   cryptSetComponent(ecdsakey->gx, buffera[3], BN_num_bits(gx));
   cryptSetComponent(ecdsakey->gy, buffera[4], BN_num_bits(gy));
   cryptSetComponent(ecdsakey->n, buffera[5], BN_num_bits(n));
   cryptSetComponent(ecdsakey->h, buffera[6], BN_num_bits(h));
   cryptSetComponent(ecdsakey->qx, buffera[7], BN_num_bits(qx));
   cryptSetComponent(ecdsakey->qy, buffera[8], BN_num_bits(qy));
cleanup:
   for (int i = 0; i < 9; i++)
      if (buffera[i])
         free(buffera[i]);
   if (ctx) BN_CTX_free(ctx);
   if (p) BN_free(p);
   if (a) BN_free(a);
   if (b) BN_free(b);
   if (gx) BN_free(gx);
   if (gy) BN_free(gy);
   if (n) BN_free(n);
   if (h) BN_free(h);
   if (qx) BN_free(qx);
   if (qy) BN_free(qy);
   return failflag;
}

/** Get a copy of the public key as an SSL format object.
 * 
 * @param[out] pubkey
 * A pointer to a place to put a pointer to the new SSL format ECDSA public key
 * which is a copy of this key. The responsibility for freeing the new
 * SSL format ECDSA public key is transferred to the caller.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int ECDSAPublicKey::getSSLKeyComponents(EC_KEY **pubkey)
{
   *pubkey = EC_KEY_new();
   if ((*pubkey) == NULL)
   {
      errprint("EC_KEY_new failed");
      return -1;
   }
   return copyECDSASSLPublicKey(*pubkey, key);
}

/** Get a copy of the public key in PEM encoded PKCS #8 text format.
 * 
 * @param[out] pemtext
 * A pointer to a place to put a pointer to the exported key in PEM
 * encoded PKCS #8 text format. The responsibility for freeing the new
 * string is transferred to the caller.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int ECDSAPublicKey::getPEMKeyData(char **pemtext)
{
   char *data;
   int damt;
   char *dbuf = NULL;
   int retval = 0;
   EVP_PKEY *epkey = EVP_PKEY_new();
   EVP_PKEY_set1_EC_KEY(epkey, key);
   BIO *bo = BIO_new(BIO_s_mem());
   if (!PEM_write_bio_PUBKEY(bo, epkey))
   {
      retval = 1;
      errprint("could not export public key");
      fSSLCheckError();
      goto bioshutdown;
   }
   fSSLCheckError();
   damt = BIO_get_mem_data(bo, &data);
   *pemtext = (char *)malloc(sizeof(char)*(damt+1));
   memcpy(*pemtext, data, sizeof(char)*(damt+1));
   (*pemtext)[damt] = '\0';
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   return retval;
}

/** Export the ECDSA public key as the binary encoding of the SSH key format
 * in the authorized_keys file.
 * 
 * @param[out] blob
 * A pointer to the place where the pointer to the binary key blob is to be
 * put. The responsibility for freeing the new buffer is transferred to the
 * caller.
 * 
 * @param[out] reallength
 * A pointer to where the write the length of the blob in bytes.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int ECDSAPublicKey::getSSHKeyBlob(char **blob, int *reallength)
{
   unsigned char keyheader[23] = {
      0x00,0x00,0x00,0x13, 'e', 'c', 'd', 's', 'a', '-',
       's', 'h', 'a', '2', '-', 'n', 'i', 's', 't', 'p',
       '2', '5', '6' };
   unsigned char curvename[12] = {
      0x00,0x00,0x00,0x08, 'n', 'i', 's', 't', 'p', '2',
       '5', '6' };
   char realnum[4];
   int nid = EC_GROUP_get_curve_name(EC_KEY_get0_group(key));
   switch (nid)
   {
      default:
      case NID_X9_62_prime256v1: sprintf(realnum, "256"); break;
      case NID_secp384r1: sprintf(realnum, "384"); break;
      case NID_secp521r1: sprintf(realnum, "521"); break;
   }
   memcpy(&(keyheader[20]), realnum, sizeof(char)*3);
   memcpy(&(curvename[9]), realnum, sizeof(char)*3);
   int enclen = 23 + 12;
   size_t enc_point_len = EC_POINT_point2oct(EC_KEY_get0_group(key),
                                             EC_KEY_get0_public_key(key),
                                             POINT_CONVERSION_UNCOMPRESSED,
                                             NULL, 0, NULL);
   unsigned char *enc_point_buf = (unsigned char *)malloc(sizeof(unsigned char)*enc_point_len);
   size_t enc_point_len2 = EC_POINT_point2oct(EC_KEY_get0_group(key),
                                              EC_KEY_get0_public_key(key),
                                              POINT_CONVERSION_UNCOMPRESSED,
                                              enc_point_buf, enc_point_len, NULL);
   enclen += (4 + enc_point_len);
   char *encode = (char *)malloc(sizeof(char)*enclen);
   memcpy(encode, keyheader, sizeof(char)*23);
   memcpy(&(encode[23]), curvename, sizeof(char)*12);
   encode[23 + 12 + 0] = (unsigned char)(enc_point_len >> 24);
   encode[23 + 12 + 1] = (unsigned char)(enc_point_len >> 16);
   encode[23 + 12 + 2] = (unsigned char)(enc_point_len >>  8);
   encode[23 + 12 + 3] = (unsigned char)(enc_point_len      );
   memcpy(&(encode[23 + 12 + 4]), enc_point_buf, sizeof(char)*enc_point_len);
   if (reallength) *reallength = 23 + 12 + 4 + enc_point_len;
   if (blob) *blob = encode;
   return 0;
}

/** Export the ECDSA public key as a line of text that can be added to the
 * SSH authorized_keys file.
 * 
 * @param[out] textblob
 * A pointer to the place where the pointer to the text-encoded labelled key
 * blob is to be put. This is the final line that can be added to
 * authorized_keys. The responsibility for freeing the new
 * string is transferred to the caller.
 * 
 * @param keylabel
 * The label to go on the key on the standardised line of text.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int ECDSAPublicKey::getSSHKeyBlob(char **textblob, const char *keylabel)
{
   int enclen = 0;
   char *encode;
   if (getSSHKeyBlob(&encode, &enclen) != 0) return -1;
   char *base64string = BaseConvert::convertBytesToBase64(encode, enclen);
   free(encode);
   char *finalstring = StringUtils::varprintf("ssh-dss %s %s", base64string, keylabel);
   free(base64string);
   if (textblob) *textblob = finalstring;
   return 0;
}

/** Compare the key to another PublicKey, and see if they are equal.
 * 
 * @param pubkey
 * PublicKey that is to be tested to see if these are equal.
 * 
 * @return
 * 1 if they match, 0 if they do not, and negative values represent
 * errors. */
int ECDSAPublicKey::compareKeyEqual(PublicKey *pubkey)
{
   if (pubkey == NULL)
   {
      errprint("comparing against NULL key");
      return 0;
   }
   switch (pubkey->getKeyType())
   {
      case KEY_TYPE_PUBLIC_ECDSA:
      {
         EC_KEY *ecdsaKeyOther = NULL;
         ((ECDSAPublicKey *)pubkey)->getSSLKeyComponents(&ecdsaKeyOther);
         EVP_PKEY *evp1 = EVP_PKEY_new();
         EVP_PKEY *evp2 = EVP_PKEY_new();
         EVP_PKEY_set1_EC_KEY(evp1, key);
         EVP_PKEY_set1_EC_KEY(evp2, ecdsaKeyOther);
         int retval = EVP_PKEY_cmp(evp1, evp2);
         EVP_PKEY_free(evp1);
         EVP_PKEY_free(evp2);
         EC_KEY_free(ecdsaKeyOther);
         fSSLCheckError();
         switch (retval)
         {
            case -2:
               errprint("key compare operation not supported (EVP_PKEY_cmp())");
               break;
            case -1:
               errprint("tried to compare different key types (EVP_PKEY_cmp())");
               return 0;
            case 0:
               return 0;
            case 1:
               return 1;
         }
         return 0;
      }
      default:
      {
         return 0;
         break;
      }
   }
}

/** Copy an SSL format public key (or the public key part of a full private
 * key in SSL format, as the public key is simply a subset of the private
 * key data).
 * 
 * @param[out] dst
 * Pointer to an initialised blank SSL format ECDSA key to copy into.
 * 
 * @param src
 * Pointer to an existing SSL format ECDSA public key to copy from.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int ECDSAPublicKey::copyECDSASSLPublicKey(EC_KEY *dst, EC_KEY *src)
{
   if (dst == NULL) { errprint("NULL public key destination"); return -1; }
   if (src == NULL) { errprint("NULL public key source"); return -1; }
   if (EC_KEY_set_group(dst, EC_KEY_get0_group(src)) == 0)
   {
      errprint("Failed to copy key group");
      return -1;
   }
   if (EC_KEY_set_public_key(dst, EC_KEY_get0_public_key(src)) == 0)
   {
      errprint("Failed to copy public key");
      return -1;
   }
   return 0;
}

/** Import an ECDSA public key from some PEM format PKCS #8 text.
 * 
 * @param text
 * The PEM encoded PKCS #8 text that represents an ECDSA public key, from
 * which this key is to be constructed.
 * 
 * @param name
 * The name to be attached to the key.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int ECDSAPublicKey::importKeyFromASCII(const char *text, const char *name)
{
   if (key) EC_KEY_free(key);
   int retval = 0;
   BIO *bo = BIO_new_mem_buf((void *)text, -1);
   EVP_PKEY *epkey = PEM_read_bio_PUBKEY(bo, NULL, 0, NULL);
   fSSLCheckError();
   if (!epkey)
   {
      retval = 1;
      errprint("could not import public key");
      key = NULL;
      goto bioshutdown;
   }
   key = EVP_PKEY_get1_EC_KEY(epkey);
   if (key == NULL)
   {
      retval = -1;
      errprint("failed to get imported public key");
   }
bioshutdown:
   BIO_set_close(bo, BIO_CLOSE);
   BIO_free(bo);
   EVP_PKEY_free(epkey);
   setKeyName(name);
   return retval;
}

/** Import and name an ECDSA public key from an ECDSA public key in encrypted
 * PEM format PKCS #8 (the encryption is UNIMPLEMENTED, this is plain text!).
 * 
 * @param text
 * The PEM encoded PKCS #8 text that represents an ECDSA public key, from
 * which this key is to be constructed.
 * 
 * @param password
 * The password to decrypt the PEM encoded PKCS #8 text that
 * represents the ECDSA public key.
 * 
 * @param name
 * The name to be attached to the key.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int ECDSAPublicKey::importKeyFromASCIIWithPassword(const char *text, const char *password, const char *name)
{
   errprint("password protection of ECDSA public keys is unsupported!");
   return importKeyFromASCII(text, name);
}

/** Export an ECDSA public key to some PEM format PKCS #8 text.
 * 
 * @param[out] text
 * A pointer to a place to put a pointer to the exported key in PEM
 * encoded PKCS #8 text format. The responsibility for freeing the new
 * string is transferred to the caller.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int ECDSAPublicKey::exportKeyToASCII(char **text)
{
   return getPEMKeyData(text);
}

/** Export an ECDSA public key to some PEM format PKCS #8 text encrypted with
 * the given password (the encryption is UNIMPLEMENTED, this is plain text!).
 * 
 * @param[out] text
 * A pointer to a place to put a pointer to the exported key in encrypted
 * PEM encoded PKCS #8 text format. The responsibility for freeing the new
 * string is transferred to the caller.
 * 
 * @param password
 * The password to encrypt and decrypt the key with.
 * 
 * @return
 * Returns zero on success, non-zero on failure. */
int ECDSAPublicKey::exportKeyToASCIIWithPassword(char **text, const char *password)
{
   errprint("password protection of ECDSA public keys is unsupported!");
   return getPEMKeyData(text);
}
